import numpy as np
import matplotlib.pyplot as plt

nx, ny = (6, 6)
x = np.linspace(0, 1, nx)
y = np.linspace(0, 1, ny)

xv, yv = np.meshgrid(x, y)

# t = [0.5, 0.5, 1]

xv = xv.reshape((-1))
yv = yv.reshape((-1))
augmented = np.full_like(xv, 1)
xy = np.vstack((xv, yv, augmented))
# print(xy.T)

# [1, 0]
# [0, 1]
r = np.array([[1, 0],
              [0, 1]])

# [0, 0]
t = np.array([0, 0])

rt = np.zeros((3, 3), dtype=np.float32)
rt[:2, :2] = r
rt[2:, :2] = t
rt[:, 2:]  = [[0],[0],[1]]

rxy = np.array([(p @ rt) for p in xy.T]).T

trxy = rxy.T
k1, k2, k3, p1, p2 = (-0.108757, -0.012452, 0.001602, -0.013687, 0.0000000)
cx = cy = 0.5
for i in range(len(trxy)):
    px, py, pw = trxy[i]
    r2 = (px - 0.5)**2 + (py - 0.5)**2
    r4 = r2**2
    r6 = r2*r4
    rx = px * (1 + k1 * r2 + k2 * r4 + k3 * r6) + (2 * p1 * py + (p2*(r2 + 2*px**2)))
    ry = py * (1 + k1 * r2 + k2 * r4 + k3 * r6) + ((p1*(r2 + 2*py**2)) + 2 * p2 * px)
    trxy[i] = [rx, ry, pw]

rxy = trxy.T

plt.plot(rxy[0], rxy[1])
plt.show()
